from functools import wraps
from flask import jsonify, request
import os
import logging


def token_required(f):
    """
    Token 验证装饰器
    :param f:
    :return:
    """
    @wraps(f)
    def decorated_function(*args, **kwargs):
        env_token = os.getenv('TOKEN', 'PXhNVVgMDdWKCrvPHOlwGseVMBscEXJKmAcrxltVQtTeNJQrMBMEIqUWSKJulDaFUUgy')

        # 支持 JSON 和 Header 两种 token 传递方式
        token = None
        if request.is_json:
            data = request.get_json(silent=True) or {}
            token = data.get('token')
        else:
            auth_header = request.headers.get('Authorization')
            if auth_header and auth_header.startswith('Bearer '):
                token = auth_header.split(' ')[1]

        if not token or token != env_token:
            logging.warning(f"Invalid token attempt: {token}")
            return jsonify({'error': 'Invalid Token'}), 401

        return f(*args, **kwargs)
    return decorated_function
